package com.yanggu.flink.state_process_api.read

import cn.hutool.core.date.DateUtil
import org.apache.flink.api.common.ExecutionConfig
import org.apache.flink.api.common.state.{ValueState, ValueStateDescriptor}
import org.apache.flink.api.scala.typeutils.Types
import org.apache.flink.configuration.Configuration
import org.apache.flink.state.api.SavepointReader
import org.apache.flink.state.api.functions.{KeyedStateReaderFunction, WindowReaderFunction}
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment
import org.apache.flink.streaming.api.windowing.assigners.WindowAssigner
import org.apache.flink.streaming.api.windowing.triggers.{CountTrigger, PurgingTrigger}
import org.apache.flink.streaming.api.windowing.windows.GlobalWindow
import org.apache.flink.util.Collector

import java.lang
import java.util.{Collections, Date}

object StateRead {

  def main(args: Array[String]): Unit = {

    val env = StreamExecutionEnvironment.getExecutionEnvironment
    env.setParallelism(1)
    val savepoint = SavepointReader.read(env, "file:/E:/tmp/flink/checkpoint/4e60c0d0fea42603ba1d68eb61233af4/chk-1")

    //测试operator list state
    val uid = "test-save-read"
    savepoint.readListState(uid, "test-list-state", Types.INT)
      .print("test-list-state")

    //测试operator union list state
    savepoint.readUnionState(uid, "test-union-list-state", Types.INT)
      .print("test-union-list-state")

    //测试broadcast state
    savepoint.readBroadcastState(uid, "test-broadcast-state", Types.STRING, Types.STRING)
      .print("test-broadcast-state")

    //测试keyed state
    savepoint.readKeyedState(uid, new KeyedStateReaderFunction[String, String] {

      //这种scala懒加载的写法是不行的, 必须在open方法中进行注册, getRuntimeContext.getState(...)方法中会将descriptor添加到本地的list中, 进而查找到状态中所有的key
      //RunTimeContext的实现是SavepointRuntimeContext。在InputFormatSourceFunction的89行format.reachedEnd, 程序会判断是不是有key, 如果没有提前注册状态, 也就没有key, 是不会执行readKey方法的
      //一个key调用一次readKey方法
      //private lazy val valueState = getRuntimeContext.getState(new ValueStateDescriptor[Int]("keyed-state-test", classOf[Int], 0))
      private var valueState: ValueState[Int] = _

      override def open(parameters: Configuration): Unit = {
        valueState = getRuntimeContext.getState(new ValueStateDescriptor[Int]("keyed-state-test", classOf[Int]))
      }

      override def readKey(key: String, ctx: KeyedStateReaderFunction.Context, out: Collector[String]): Unit = {
        out.collect(s"key: $key, value-state: ${valueState.value()}")
      }
    }).print("test-keyed-state")

    //测试window state
    savepoint.window(new WindowAssigner[Object, GlobalWindow] {

      override def assignWindows(element: Object, timestamp: Long, context: WindowAssigner.WindowAssignerContext) = Collections.singletonList(GlobalWindow.get())

      override def getDefaultTrigger(env: StreamExecutionEnvironment) = PurgingTrigger.of[Object, GlobalWindow](CountTrigger.of[GlobalWindow](5L))

      override def getWindowSerializer(executionConfig: ExecutionConfig) = new GlobalWindow.Serializer

      override def isEventTime = false
    })
      .process("test-save-read-window", (key: String, context: WindowReaderFunction.Context[GlobalWindow], elements: lang.Iterable[String], out: Collector[String]) => {
        val window = context.window()
        out.collect(s"${DateUtil.formatDateTime(new Date(window.maxTimestamp()))} elements: ${elements.toString}")
      }, Types.STRING, Types.STRING, Types.STRING)
      .print("test-count-window-state")

    env.execute()

  }

}
